require(stringi)
require(runjags)
require(coda)
require(doParallel)

party.position.sample <- array(NA, c(46, 12, 61, 100))  # parties x committees x years x bootstrap draws
party.position.data <- c()
for (i in 1:61) {
  raw.data <- readRDS(paste0("Wordfish_bootstrap/positions_", 1958 + i, "_boot.rds"))
  party.position.data <- rbind(party.position.data, 
                               data.frame(party = rownames(raw.data), 
                                          year = 1958 + i, time = i, 
                                          raw.data[, , 1]))
  party.position.sample[, , i, ] <- raw.data
}

party.position.data <- subset(party.position.data, 
                              apply(party.position.data, 1, function(x) sum(is.na(x)) < 12))

party.position.data$party.id <- as.numeric(factor(party.position.data$party))  # party ID

save(party.position.data, file = "Wordfish_bootstrap/bootstrap_array.Rdata")

l <- nrow(party.position.data)  # number of observations (party-year cells with non-missing data)
n <- max(party.position.data$party.id)  # number of parties
m <- sum(stri_detect_regex(colnames(party.position.data), "[一-龠]+"))  # number of committees
t <- max(party.position.data$time)  # number of years

party.begin <- party.end <- rep(NA, n)  # first and last observed year for each party
for (i in 1:n) {
  party.begin[i] <- min(party.position.data$time[party.position.data$party.id == i])
  party.end[i] <- max(party.position.data$time[party.position.data$party.id == i])
}
committee.begin <- committee.end <- rep(NA, m)  # first and last observed year for each committee
for (j in 1:m) {
  temp <- subset(party.position.data, ! is.na(party.position.data[, j + 3]))
  committee.begin[j] <- min(temp$time)
  committee.end[j] <- max(temp$time)
}

party.valid.columns <- rep(FALSE, n * t)
for (i in 1:n) {
  for (j in 1:t) {
    party.valid.columns[n * (j - 1) + i] <- party.begin[i] <= j & j <= party.end[i]
  }
}

committee.valid.columns <- rep(FALSE, m * t)
for (i in 1:m) {
  for (j in 1:t) {
    committee.valid.columns[m * (j - 1) + i] <- committee.begin[i] <= j & j <= committee.end[i]
  }
}

# Note: The party IDs are determined by factor(party); the initial values below assume a fixed ID mapping.
initial.fun <- function(seed1, seed2) {
  set.seed(seed1)
  theta <- matrix(NA, n, t)
  theta[, 1] <- rnorm(n)
  theta[2, 1] <- -2  # set a negative initial value for the JCP (HOR)
  theta[6, 1] <- 2  # set a positive initial value for the LDP (HOR)
  beta <- matrix(NA, m, t)
  beta[, 1] <- runif(m, 0, 1)
  inv.phi2 <- runif(m, 0.5, 1.5)
  inv.xi2 <- runif(1, 0.5, 1.5)
  kappa <- runif(1, 0.1, 0.2)
  lambda <- runif(1, 0.1, 0.2)
  list(beta.raw = beta, theta.raw = theta, 
       inv.phi2 = inv.phi2, inv.xi2 = inv.xi2, 
       kappa = kappa, lambda = lambda, 
       .RNG.name = "base::Wichmann-Hill", .RNG.seed = seed2)
}
initial.values <- list()
for (i in 1:3) {
  initial.values[[i]] <- initial.fun(100 * i, 1000 * i)
}

# Note: The object name contains a typo ("Wordshaol"), kept as-is to match the original analysis.
for (i in 1:100) {
  dynamic.Wordshaol.result <- run.jags(model = "1-b_dynamic_factor_analysis.R", 
                                       monitor = c("theta", "beta", "phi", "tau", "omega"), 
                                       data = list(y = party.position.sample[, , , i], n = n, m = m, t = t), 
                                       inits = initial.values, 
                                       n.chains = 3, adapt = 1000, burnin = 5000, sample = 5000, modules = "glm", 
                                       summarise = FALSE, thin = 100, method = "parallel")
  save(dynamic.Wordshaol.result, file = paste0("Wordshoal/Wordshoal_", i, ".Rdata"))
}

convergence.check <- rep(NA, 100)
for (i in 1:100) {
  load(paste0("Wordshoal/Wordshoal_", i, ".Rdata"))
  convergence.check[i] <- sum(gelman.diag(dynamic.Wordshaol.result, multivariate = FALSE)$psrf[, 1] > 1.1)
  cat("Iteration", i, "has been finished.\n")
}
convergence.check

# total posterior draws: 100 bootstrap datasets × 3 chains × 500 thinned draws = 150,000
posterior.draws.party.raw <- matrix(NA, 150000, sum(party.valid.columns))
posterior.draws.committee.raw <- matrix(NA, 150000, sum(committee.valid.columns))
load("Wordshoal/Wordshoal_1.Rdata")
colnames(posterior.draws.party.raw) <- colnames(dynamic.Wordshaol.result$mcmc[[1]][, c(1:(n * t))[party.valid.columns]])
colnames(posterior.draws.committee.raw) <- colnames(dynamic.Wordshaol.result$mcmc[[1]][, c((n * t + 1):((n + m) * t))[committee.valid.columns]])
for (i in 1:100) {
  load(paste0("Wordshoal/Wordshoal_", i, ".Rdata"))
  posterior.draws.party.raw[(1500 * (i - 1) + 1):(1500 * i), ] <- rbind(dynamic.Wordshaol.result$mcmc[[1]][seq(1, 5000, 10), c(1:(n * t))[party.valid.columns]], 
                                                                        dynamic.Wordshaol.result$mcmc[[2]][seq(1, 5000, 10), c(1:(n * t))[party.valid.columns]], 
                                                                        dynamic.Wordshaol.result$mcmc[[3]][seq(1, 5000, 10), c(1:(n * t))[party.valid.columns]])
  posterior.draws.committee.raw[(1500 * (i - 1) + 1):(1500 * i), ] <- rbind(dynamic.Wordshaol.result$mcmc[[1]][seq(1, 5000, 10), c((n * t + 1):((n + m) * t))[committee.valid.columns]], 
                                                                            dynamic.Wordshaol.result$mcmc[[2]][seq(1, 5000, 10), c((n * t + 1):((n + m) * t))[committee.valid.columns]], 
                                                                            dynamic.Wordshaol.result$mcmc[[3]][seq(1, 5000, 10), c((n * t + 1):((n + m) * t))[committee.valid.columns]])
}

# post-processing: standardize theta within each draw (anchored at party.begin) and rescale beta accordingly
iteration.mean <- rowMeans(posterior.draws.party.raw[, paste0("theta[", 1:length(party.begin), ",", party.begin, "]")])
iteration.sd <- apply(posterior.draws.party.raw[, paste0("theta[", 1:length(party.begin), ",", party.begin, "]")], 1, sd)

cluster <- makeCluster(10)
registerDoParallel(cluster)
postprocessing.out <- foreach(i = 1:150000, .combine = rbind) %dopar% {
  c(std.theta = (posterior.draws.party.raw[i, ] - iteration.mean[i]) / iteration.sd[i], 
    std.beta = iteration.sd[i] * posterior.draws.committee.raw[i, ])
}
stopCluster(cluster)

save(postprocessing.out, file = "Wordshoal/postprocessing_out.Rdata")
